# -*- coding: utf-8 -*-
# $Id: udineGenUtils2.py 3 2010-01-06 08:25:59Z leo.monash.uni $
import math, string, glob
import pickle, sys
from math import ceil
from random import seed, random, choice, randrange, getstate, setstate
import networkx

class TimetablingInstance:
  "A class storing properties of a (randomly generated) instance of a course timetabling problem"
  # The connected components from all the conflict graphs from curricula 
  # from the 21 itc2007 instances. 
  itc07Graph, itc07enrollment = pickle.load(open('combinedITC07graphs.pickle'))
  itc07components = networkx.connected_component_subgraphs(itc07Graph)
  
  def __init__(self, events, occupancy,nameConvention="Name: Random_%iEvents_Occupancy%i\n"):
    #seed(1) # uncomment for debugging.
    self.pEvents = events
    self.pOccupancy = float(occupancy)/100
    self.nameConvention = nameConvention

    self.pDays = 5
    self.pPeriodsPerDay = 5
    self.pPeriods = self.pDays * self.pPeriodsPerDay 

    self.pMinEventsPerCourse = 2
    self.pMaxEventsPerCourse = 5
    self.pCourses = 0            # reset in genCourses!
    self.pMinDaysFactor = 1.125
    self.pMaxRoomCourseRatio = 1
    self.pMinRoomCourseRatio = 0.7
    self.enrollment = {}

    self.pTeachers = int(self.pEvents / 6)
    self.pMinCoursesPerTeacher = 1
    self.pMaxCoursesPerTeacher = 3

    self.pCurricula = int(self.pEvents / 10)
    self.pMinCoursesPerCurriculum = 2
    self.pMaxCoursesPerCurriculum = 6
    # curriculaGraph and associated data are for sampling.
    self.curriculaGraph = networkx.Graph() 
    self.curriculaGraphComponents = []

    self.pRooms = max(int(ceil(self.pEvents / self.pPeriods / self.pOccupancy)), 1)
    self.pDistinctRoomSizes = max(int(ceil(0.6 * self.pRooms)), 1)
    self.pSmallestClassroom = 40
    self.pLargestClassroom = 350

    self.pConstraints = int(self.pEvents * 1.2)

    self.intRoomSizes = []
    self.intCourseIds = []
    

  def genRooms(self):
    outRooms = "\nROOMS:\n"
    distinct = []
    for i in range(self.pDistinctRoomSizes):
      distinct.append(randrange(self.pSmallestClassroom, self.pLargestClassroom, 5))
    complete = distinct[:]
    for j in range(self.pRooms - self.pDistinctRoomSizes):
      complete.append(choice(distinct))
    complete.sort()
    intRoomSizes = complete[:] 
    for room, capacity in zip(range(self.pRooms), complete):
      outRooms += "R%02i\t%i\n" % (room + 1, capacity)
    return outRooms
    
  def genCourses(self):
    outCourses = "\nCOURSES:\n"
    eventsCounts = []
    minDaysSpecs = []
    taughtEvents = 0
    while taughtEvents < self.pEvents:
      ev = min(randrange(self.pMinEventsPerCourse, self.pMaxEventsPerCourse), self.pEvents - taughtEvents)
      eventsCounts.append(ev)
      minDaysSpecs.append(int(ev/self.pMinDaysFactor))
      taughtEvents += ev
    self.pCourses = len(eventsCounts)
    taught = self.pMinCoursesPerTeacher * range(self.pTeachers)
    teachers = set(range(self.pTeachers))
    while len(taught) < self.pCourses:
      teacher = teachers.pop()
      for i in range(max(self.pMaxCoursesPerTeacher - self.pMinCoursesPerTeacher, len(taught) - self.pCourses)):
        taught.append(teacher)
    students = [randrange(int(self.pSmallestClassroom/2), self.pLargestClassroom) for j in range(self.pCourses)]
    for course, teacher, eventsCount, minDays, capacity in zip(range(self.pCourses), taught, eventsCounts, minDaysSpecs, students):
      id = "Course%03i" % (course + 1)
      self.intCourseIds.append(id)
      # keep track of enrollment to use for sampling curricula later.
      self.enrollment[id] = capacity
      outCourses += "%s t%02i %i %i %i\n" % (id, teacher, eventsCount, minDays, capacity)
    return outCourses

  def genUnavailability(self):
    outUnavail = "\nUNAVAILABILITY_CONSTRAINTS:\n"
    unavail = []
    while len(unavail) < self.pConstraints:
      course = choice(self.intCourseIds)
      day = randrange(0, self.pDays)
      period = randrange(0, self.pPeriodsPerDay)
      if random() > 0.5:
        for p in range(period, self.pPeriodsPerDay): unavail.append((course, day, p))
      else:
        for p in range(0, period+1): unavail.append((course, day, p))
    unavail = sorted(list(set(unavail))) # May result in the number of constraints being cut, so...
    self.pConstraints = len(unavail) # ... update the number of constraints.
    for course, day, period in unavail:
      outUnavail += "%s %i %i\n" % (course, day, period)
    return outUnavail

  def genCurricula(self):

    def choiceWithWeights(elemFreqTuples):
      """Like choice but takes into account weights."""
      X = random()*sum([f for (i,f) in elemFreqTuples])
      s = 0.
      for (i,f) in elemFreqTuples:
        s += f
        if X < s: break
      return i

    def sampleGraph():
      """Produce a graph based on the itc2007 connected components 
      where the number of nodes matches the number of courses selected above."""
      # Only interested in components longer than shortest curriculum.
      largerThanShortestCurriculum = [c for c in self.itc07components if len(c.nodes()) >= self.pMinCoursesPerCurriculum]
      components, nCourses, = [], 0
      while(nCourses < self.pCourses):
        lastComponent = choice(largerThanShortestCurriculum)
        components.append(lastComponent)
        nCourses = sum([len(c) for c in components])
      # find the component which has to be modified the least to fit.
      # Try to keep the large components as intact as possible
      # since they are more likely to constrain the problem.
      if nCourses > self.pCourses:
        excess = nCourses - self.pCourses
        cutSize, toCut = sorted([(len(c)-excess,c) for c in components if len(c)-excess >=0])[0]
        del components[components.index(toCut)]
        # if we've cut too many classes by removing toCut, add back a connected subgraph
        if cutSize:
          components.append(toCut.subgraph(curriculumFromWeighedWalk(cutSize,toCut,self.itc07enrollment)))
      # rename the nodes with enrollments in the same order. 
      fullGraph = self.itc07Graph.subgraph(sum([c.nodes() for c in components],[]))
      # print fullGraph.nodes()
      itcNodesBySize = sorted(fullGraph.nodes(),key=lambda x: self.itc07enrollment[x])
      thisInstanceNodesBySize = sorted(self.enrollment,key = lambda x: self.enrollment[x])
      translation = dict(zip(itcNodesBySize,thisInstanceNodesBySize))
      self.curriculaGraph = networkx.Graph()
      for (old,new) in translation.iteritems():
        self.curriculaGraph.add_node(new)
        for v in fullGraph.neighbors(old): self.curriculaGraph.add_edge(new,translation[v])
      # compute this helpful quantity only once
      self.curriculaGraphComponents = networkx.connected_component_subgraphs(self.curriculaGraph)
      
    def curriculumFromWeighedWalk(length, graph=None,weights=None,initialNode=None):
      """Returns a random set of courses of length 'length' connected through a graph.
      The list is obtained from a random walk where the probability of a step is
      obtained by the relative weights of nodes. Thus this procedure can cycle for ever.
      'graph' is a networkx.Graph. 
      'weights' is a dict with entries for each node in 'graph'
      """
      # random walks can take a long time to complete if 
      # some courses have very low enrollment and length is close to the size of the component.
      # The following lines help in those cases.
      if initialNode and initialNode not in graph.nodes(): 
        raise KeyError, "No such node %s. Nodes: %s" % (initialNode,graph.nodes())
      if length == len(graph): seq = set(graph.nodes())
      else:
        if initialNode: last = initialNode
        else: last = choiceWithWeights([(v,weights[v]) for v in graph.nodes()])
        seq = set([last])
        while len(seq) < length:
          last = choiceWithWeights([(v,weights[v]) for v in graph.neighbors(last)])
          seq.add(last)
      return seq

    outCurr = "\nCURRICULA:\n"
    curr = []
    # Initialize sample graph. Graph must have at least one component large enough 
    # to generate the largest curriculum.
    sampleGraph()
    # first make sure every course is used at least once
    unusedCourses = set(self.curriculaGraph.nodes())
    nCurricula = 0
    while unusedCourses:
      start = choiceWithWeights([(c,self.enrollment[c]) for c in unusedCourses])
      component = [g for g in self.curriculaGraphComponents if start in g.nodes()][0]
      nCourses = randrange(min(self.pMinCoursesPerCurriculum,len(component)),min(self.pMaxCoursesPerCurriculum,len(component))+1)
      courses = curriculumFromWeighedWalk(nCourses,component,self.enrollment,start)
      curr.append((nCurricula,courses))
      unusedCourses.difference_update(courses)
      nCurricula += 1
    # now if there are any curricula left add them.
    for i in range(nCurricula,self.pCurricula):
      nCurricula += 1
      nCourses = randrange(self.pMinCoursesPerCurriculum,self.pMaxCoursesPerCurriculum+1)
      bigEnoughComponents = [g for g in self.curriculaGraphComponents if len(g.nodes()) >= nCourses]
      component = choice(bigEnoughComponents)
      courses = curriculumFromWeighedWalk(nCourses,component,self.enrollment)
      curr.append((i, courses))
    for id, courses in curr:
      courses = sorted(list(courses))
      outCurr += "Curr%003i %i %s\n" % (id, len(courses), string.join(courses, " "))
    self.pCurricula = nCurricula
    return outCurr

  def genHeader(self,name=None):
    if name: outHeader = "Name: %s\n" % name
    else: outHeader = self.nameConvention % (self.pEvents, int(self.pOccupancy * 100))
    outHeader += "Courses: %i\n" % self.pCourses
    outHeader += "Rooms: %i\n" % self.pRooms
    outHeader += "Days: %i\n" % self.pDays
    outHeader += "Periods_per_day: %i\n" % self.pPeriodsPerDay
    outHeader += "Curricula: %i\n" % self.pCurricula
    outHeader += "Constraints: %i\n" % self.pConstraints
    return outHeader

  def genInstanceDef(self,name=None):
    rooms = self.genRooms()
    courses = self.genCourses()
    curr = self.genCurricula()
    unavail = self.genUnavailability()
    header = self.genHeader(name)
    return header + courses + rooms + curr + unavail + "\nEND."

if __name__ == '__main__':
  import sys
  #setstate(pickle.load(open('state.pickle')))
  try:
    if len(sys.argv) == 3:
      state = getstate()
      print TimetablingInstance(int(sys.argv[1]),int(sys.argv[2])).genInstanceDef()
    elif len(sys.argv) == 2:
      state,argv,e = pickle.load(open(sys.argv[1]))
      print "Handling exception: \n\n%s\n\n"
      setstate(state)
      print TimetablingInstance(int(argv[1]),int(argv[2])).genInstanceDef()
  except Exception as e:
    print e
    pickle.dump([state,sys.argv,e],open('state.%s-%9d.pickle' % sys.argv[1:],random.randint(0,999999999),'w'))
    raise

unitTest = """
#bash script:
for events in 20 40 50 70 100 125 150 200 400; do 
  for occupancy in 30 50 70 80 90; do 
    for i in `seq 1 100`; do 
      echo -n ${events} ${occupancy} ${i}", "; 
      python udineGenUtils2.py ${events} ${occupancy} > ${events}Events${occupancy}PctOccupancy-${i}.ctt; 
    done 
  done 
done
"""
